from util.pic_video import *
from util.LabelPoints2ImagebyMayaParam import *
from util.visualize_utils import *
import os
import time
import pandas as pd
import json
import shutil

type_table = {
    '轿车': 'car',
    '巴士': 'bus',
    '货车': 'truck',
    '行人':'pedestrians',
    '挂车':'trailer',
    '其他':'other',
    '工程车辆':'Engineering_vehicles',
    '交通锥':'Traffic_cone',
    '两轮车':'motorcycle'
}

def get_json_data(data_label):
    # 读取预测数据
    name = [i['obj_type'].lower() for i in data_label]
    xyz = np.array(
        [[i['psr']['position']['x'], i['psr']['position']['y'], i['psr']['position']['z']] for i in data_label])
    lwh = np.array([[i['psr']['scale']['x'], i['psr']['scale']['y'], i['psr']['scale']['z']] for i in data_label])
    r = np.array([[i['psr']['rotation']['z']] for i in data_label])
    box = np.concatenate([xyz, lwh, r], axis=-1)
    # box[:, 2] = box[:, 2] - 3.15
    velocity = np.zeros([name.__len__(), 5]).tolist()
    acceleration = np.zeros([name.__len__(), 5]).tolist()
    return name, box, velocity, acceleration

file='/home/king/workplace/data/'
alone=['data14_2021_04_03','data27_2021_04_03','data28_2021_04_03']

for alone_ in alone:
    data_file = os.path.join(file,alone_)
    print(data_file)

    # 解析整体gt文件成json写入gtlabel
    gtlabel_path=data_file + '/gt_label_json'
    if not os.path.exists(gtlabel_path):
        os.makedirs(gtlabel_path)
    ground_truth = data_file + '/ground_truth/'
    gt_files=os.listdir(ground_truth)
    with open(ground_truth+gt_files[0], 'r') as f:
        data_temp = json.load(f)
    for file_id in range(data_temp.__len__()):
        datas_temp = []
        file_name = data_temp[file_id]['3D_source'].split('/')[-1][0:-4] + '.json'
        for obj_id in range(data_temp[file_id]['框'].__len__()):
            data1 = {'psr':
                         {'position':
                              data_temp[file_id]['框'][obj_id]['center'],
                          'scale': {'x': data_temp[file_id]['框'][obj_id]['width'],
                                    'y': data_temp[file_id]['框'][obj_id]['height'],
                                    'z': data_temp[file_id]['框'][obj_id]['depth']},
                          'rotation': data_temp[file_id]['框'][obj_id]['rotation']},
                     'obj_type': type_table[data_temp[file_id]['框'][obj_id]['attributes']['类型'][0]],
                     'obj_id': data_temp[file_id]['框'][obj_id]['id']
                     }
            datas_temp.append(data1)
        name, box, velocity, acceleration = get_json_data(datas_temp)
        with open(gtlabel_path+ '/' + file_name, 'w') as f:
            json.dump(datas_temp, f)

    # 复制data内.bin文件到lidar
    lidar_path=data_file+'/lidar/'
    if not os.path.exists(lidar_path):
        shutil.copytree(data_file+'/data_bin/',lidar_path)

    # # 复制pic_040302_release/p0文件到image
    # image_path=data_file+'/image/'
    # if not os.path.exists(image_path):
    #     shutil.copytree(data_file+'/pic_040302_release/p0/',image_path)